# Dickstein, Ho, and Mark (2023)
# Create the relevant heat maps and other outputs related to varying parameters analysis

# Preliminaries
setwd("../../../../library")
source("PreliminariesCode.R")
library(foreign)
library(cowplot)
library(grid)
library(gridExtra)

data_folder <- paste0(project_folder, "/data")

build_folder <- paste0(project_folder, "/analysis/tablesandfigures/build/heatmaps")
dir.create(build_folder, recursive = T)

release_folder <- paste0(project_folder, "/analysis/tablesandfigures/release/heatmaps")
dir.create(release_folder, recursive = T)

table <- fread(paste0(data_folder, "/outcome_table_hm.csv"))

psi_upperbound <- .3

# Creating "Change" variables: 
vars <- unique(gsub("im.pre", "", names(table)[which(substr(names(table), 1, 6) == "im.pre")]))
for(i in seq_along(vars)){
  table[, eval(parse(text = paste0("im.change", vars[i], " := im.post", vars[i], " - im.pre", vars[i]))), ]
  table[, eval(parse(text = paste0("sg.change", vars[i], " := sg.post", vars[i], " - sg.pre", vars[i]))), ]
}

# Find the baseline parameter gamble value
gamble_interpret <- function(b3){1000 * exp(-b3) * log(2 - exp(-exp(b3)/10))}
b3_im <- as.numeric(unlist(table[307, "im.beta3"]))
b3_sg <- as.numeric(unlist(table[307, "sg.beta3"]))
im_baseline_gamble <- gamble_interpret(b3_im)
sg_baseline_gamble <- gamble_interpret(b3_sg)

# Creating param variables:
table[, im.omega := exp(im.beta2), ]
table[, sg.omega := exp(sg.beta2), ]
table[, im.psi := exp(im.beta3), ]
table[, sg.psi := exp(sg.beta3), ]
table[, im.which_gamble := gamble_interpret(im.beta3)]
table[, sg.which_gamble := gamble_interpret(sg.beta3)]

# Creating plots:
figure_params_table <- data.table(
  x_axis = c(
    "im.omega", "im.omega", "im.lambdashifter_ind", "im.lambdashifter_ind"), 
  y_axis = c(
    "sg.omega", "sg.omega", "sg.lambdashifter_sg", "sg.lambdashifter_sg"),
  z_axis = c(
    "im.change.cs_2", "sg.change.cs_2", "im.change.cs_2", "sg.change.cs_2"),
  title = c(
    "Change in IM Consumer Surplus", "Change in SG Consumer Surplus", "Change in IM Consumer Surplus", "Change in SG Consumer Surplus"), 
  x_lab = c(
    "Individual market moral hazard", "Individual market moral hazard", "Individual market lambda multiplier", "Individual market lambda multiplier"),
  y_lab = c(
    "Small group market moral hazard", "Small group market moral hazard", "Small group market lambda multiplier", "Small group market lambda multiplier"),
  z_lab = c(
    "change in consumer surplus", "change in consumer surplus", "change in consumer surplus", "change in consumer surplus"),
  start_obs = c(1, 1, 234, 234),
  end_obs = c(81, 81, 333, 333)
)

# lotting each figure
for(j in 1:nrow(figure_params_table)){
  data_temp <- table[figure_params_table$start_obs[j]:figure_params_table$end_obs[j]]
  p <- ggplot(
    data = data_temp,
    aes(
      x = eval(parse(text = figure_params_table$x_axis[j])),
      y = eval(parse(text = figure_params_table$y_axis[j]))
    ))
  
  p <- p + geom_tile(aes(fill = eval(parse(text = figure_params_table$z_axis[j]))))
  p <- p + scale_fill_distiller(palette = "Greys")
  p <- p + theme_bw()
  p <- p + labs(
    x = figure_params_table$x_lab[j],
    y = figure_params_table$y_lab[j], 
    fill = figure_params_table$z_lab[j])
  ## Omega graphs
  if(j %in% c(1,2)){
    p <- p + geom_vline(xintercept = exp(as.numeric(unlist(table[307, "im.beta2"]))), linetype=2)
    p <- p + geom_hline(yintercept = exp(as.numeric(unlist(table[307, "sg.beta2"]))), linetype=2)
  ## Psi graphs
  } else {
    p <- p + geom_vline(xintercept = 1, linetype=2) 
    p <- p + geom_hline(yintercept = 1, linetype=2)
    p <- p + geom_point(aes(x=.766,y=1), shape = 21, colour = "black", fill = "white", size = 5)
  }
  p <- p + theme(legend.position="bottom", legend.justification = "center", text = element_text(size = 14))
  legend <- get_legend(p)
  p_leg <- arrangeGrob(legend, nrow = 1)
  p <- p + theme(legend.position = "none", text = element_text(size = 14))
  ggsave(plot = p, filename = paste0(
    release_folder, "/",
    gsub (" ", "", figure_params_table$x_lab[j], fixed = TRUE), 
    gsub (" ", "", figure_params_table$title[j], fixed = TRUE),
    ".eps"), 
    device = "eps")
      
  ggsave(plot = p_leg, filename = paste0(
    release_folder, "/",
    gsub (" ", "", figure_params_table$x_lab[j], fixed = TRUE), 
    gsub (" ", "", figure_params_table$title[j], fixed = TRUE),
    "_legend",
    ".eps"), 
    width = 6.17277700913242, height = 1.48732367579909,
    device = "eps")
}

# Export table as .dta file to make premarkup graphs in STATA
## Keep only relevant rows and add labels
table_mkup1 <- table[182:207, ]
table_mkup2 <- table[208:233, ]
table_mkup1[, premrkup := 0]
table_mkup2[, premrkup := .25]
table_mrkup <- rbind(table_mkup1, table_mkup2)

## Saving
write.dta(table_mrkup, file = paste0(build_folder, "/premarkup_counterfactual.dta"))

# Export table as .dta file to make risk aversion line graphs in STATA
## Keep only relevant rows and add labels
table_riskav_sg_same <- table[(81+91):(81+100), ] 
im_cs_norm <- as.numeric(table[81+100, "im.post.cs_2"]) 
im_same_rows <- c(3, 13, 23, 33, 43, 53, 63, 73, 83, 93)
im_same_rows <- 81 + im_same_rows
im_cs_norm2 <- as.numeric(table[81+93, "im.post.cs_2"])
table_riskav_im_same <- table[im_same_rows,]
table_riskav_sg_same[, sg_same := 1]
table_riskav_im_same[, sg_same := 0]

## Normalize the consumer surpluses for purpose of risk aversion line graph
table_riskav_sg_same[, im.post.cs_2 := im.post.cs_2 - im_cs_norm]
table_riskav_sg_same[, im.pre.cs_2 := im.pre.cs_2 - im_cs_norm]
table_riskav_im_same[, im.post.cs_2 := im.post.cs_2 - im_cs_norm2]
table_riskav_im_same[, im.pre.cs_2 := im.pre.cs_2 - im_cs_norm2]
table_riskav <- rbind(table_riskav_sg_same, table_riskav_im_same)

## Saving
write.dta(table_riskav, file = paste0(build_folder, "/riskaversion_counterfactual.dta"))
